"""
LLM Agent配置

配置LLM Agent的接口地址、认证信息和并发参数
"""

from typing import Dict, Any
from pydantic import BaseModel, Field


class AgentConfig(BaseModel):
    """单个Agent配置"""
    agent_id: str = Field(..., description="Agent ID")
    base_url: str = Field(..., description="API基础URL")
    api_key: str = Field(..., description="API密钥")
    timeout: int = Field(default=300, description="请求超时时间（秒）")
    max_retries: int = Field(default=3, description="最大重试次数")


class ConcurrencyConfig(BaseModel):
    """并发配置"""
    extraction_workers: int = Field(default=5, description="字段提取并发数")
    tracing_workers: int = Field(default=5, description="字段溯源并发数")
    audit_workers: int = Field(default=5, description="字段审核并发数")
    classification_workers: int = Field(default=5, description="文件分类并发数")


class LLMConfig:
    """LLM配置"""
    
    # 开发环境配置
    DEV_CONFIG = {
        "extraction_agent": AgentConfig(
            agent_id="f526d16a58db42698d574adfba7fb397",
            base_url="https://uat.agentspro.cn/openapi/agents/chat/completions/v1",
            api_key="02d121e8db8d42f39b7a26d0843e4c12.HgdrT9qnyk1k584xpT0hhseGg3Q7u9at",
            timeout=300,
            max_retries=3
        ),
        "tracing_agent": AgentConfig(
            agent_id="34f65628b11c41a8ad16c5aa32c7b458",
            base_url="https://uat.agentspro.cn/openapi/agents/chat/completions/v1",
            api_key="02d121e8db8d42f39b7a26d0843e4c12.HgdrT9qnyk1k584xpT0hhseGg3Q7u9at",
            timeout=300,
            max_retries=3
        ),
        "audit_agent": AgentConfig(
            agent_id="0c3a2c543c234c20a20447c43420ce0f",
            base_url="https://uat.agentspro.cn/openapi/agents/chat/completions/v1",
            api_key="02d121e8db8d42f39b7a26d0843e4c12.HgdrT9qnyk1k584xpT0hhseGg3Q7u9at",
            timeout=300,
            max_retries=3
        ),
        "classification_agent": AgentConfig(
            agent_id="362a81dfdf0344deb8541e232b4eac34",
            base_url="https://uat.agentspro.cn/openapi/agents/chat/completions/v1",
            api_key="02d121e8db8d42f39b7a26d0843e4c12.HgdrT9qnyk1k584xpT0hhseGg3Q7u9at",
            timeout=60,
            max_retries=3
        )
    }
    
    # 生产环境配置（待配置）
    PROD_CONFIG = {
        "extraction_agent": AgentConfig(
            agent_id="PROD_EXTRACTION_AGENT_ID",  # TODO: 替换为生产环境Agent ID
            base_url="https://api.agentspro.cn/openapi/agents/chat/completions/v1",  # TODO: 替换为生产环境URL
            api_key="PROD_API_KEY",  # TODO: 替换为生产环境API Key
            timeout=300,
            max_retries=3
        ),
        "tracing_agent": AgentConfig(
            agent_id="PROD_TRACING_AGENT_ID",  # TODO: 替换为生产环境Agent ID
            base_url="https://api.agentspro.cn/openapi/agents/chat/completions/v1",  # TODO: 替换为生产环境URL
            api_key="PROD_API_KEY",  # TODO: 替换为生产环境API Key
            timeout=300,
            max_retries=3
        ),
        "audit_agent": AgentConfig(
            agent_id="PROD_AUDIT_AGENT_ID",  # TODO: 替换为生产环境Agent ID
            base_url="https://api.agentspro.cn/openapi/agents/chat/completions/v1",  # TODO: 替换为生产环境URL
            api_key="PROD_API_KEY",  # TODO: 替换为生产环境API Key
            timeout=300,
            max_retries=3
        ),
        "classification_agent": AgentConfig(
            agent_id="PROD_CLASSIFICATION_AGENT_ID",  # TODO: 替换为生产环境Agent ID
            base_url="https://api.agentspro.cn/openapi/agents/chat/completions/v1",  # TODO: 替换为生产环境URL
            api_key="PROD_API_KEY",  # TODO: 替换为生产环境API Key
            timeout=60,
            max_retries=3
        )
    }
    
    # 并发配置
    CONCURRENCY_CONFIG = ConcurrencyConfig(
        extraction_workers=5,  # 字段提取并发数
        tracing_workers=5,     # 字段溯源并发数
        audit_workers=5,       # 字段审核并发数
        classification_workers=5  # 文件分类并发数
    )
    
    # 当前环境（dev/prod）
    CURRENT_ENV = "dev"
    
    @classmethod
    def get_extraction_agent_config(cls) -> AgentConfig:
        """获取字段提取Agent配置"""
        config = cls.DEV_CONFIG if cls.CURRENT_ENV == "dev" else cls.PROD_CONFIG
        return config["extraction_agent"]
    
    @classmethod
    def get_tracing_agent_config(cls) -> AgentConfig:
        """获取字段溯源Agent配置"""
        config = cls.DEV_CONFIG if cls.CURRENT_ENV == "dev" else cls.PROD_CONFIG
        return config["tracing_agent"]
    
    @classmethod
    def get_audit_agent_config(cls) -> AgentConfig:
        """获取字段审核Agent配置"""
        config = cls.DEV_CONFIG if cls.CURRENT_ENV == "dev" else cls.PROD_CONFIG
        return config["audit_agent"]
    
    @classmethod
    def get_classification_agent_config(cls) -> AgentConfig:
        """获取文件分类Agent配置"""
        config = cls.DEV_CONFIG if cls.CURRENT_ENV == "dev" else cls.PROD_CONFIG
        return config["classification_agent"]
    
    @classmethod
    def get_concurrency_config(cls) -> ConcurrencyConfig:
        """获取并发配置"""
        return cls.CONCURRENCY_CONFIG
    
    @classmethod
    def set_environment(cls, env: str):
        """
        设置当前环境
        
        Args:
            env: 环境名称 ('dev' 或 'prod')
        """
        if env not in ["dev", "prod"]:
            raise ValueError(f"无效的环境配置: {env}，必须是 'dev' 或 'prod'")
        cls.CURRENT_ENV = env


# 导出配置实例
llm_config = LLMConfig()

